Wraps the model with native Torch FSDP. Experimental!